import datetime
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from matplotlib.widgets import MultiCursor
from pandas.plotting import register_matplotlib_converters

import QUANTAXIS as QA
from QAStrategy.QAStrategy.qa_muti_freq_stock import QAStrategyMutiFreqStockBase


class MutiFreqStrategy(QAStrategyMutiFreqStockBase):
    '''
    多周期共振策略
    '''
    def __init__(self, username, password, frequences, code='510300', start='2019-01-01', end='2019-10-21',
                 min_trade_money=5000, show_trade=True, once_trade_money=10000,
                 strategy_id='QA_STRATEGY', init_cash=1000000, market=QA.MARKET_TYPE.STOCK_CN):
        super().__init__(username=username, password=password, code=code, start=start, end=end,
                         frequences=frequences, strategy_id=strategy_id, init_cash=init_cash)
        self.market_type = market
        self.start = datetime.datetime.strptime(start, '%Y-%m-%d')
        self.end = datetime.datetime.strptime(end, '%Y-%m-%d')
        self.trade_start_time = None
        self.show_Trade_log = show_trade
        self.min_trade_money = min_trade_money
        self.once_trade_money = once_trade_money
        self.min2day_coef = 1

    # sim account on bar
    def on_bar(self, bars):
        current_dt = bars.index.values[0][0]
        current_bar = bars.xs(current_dt, level=0)
        JC = 1
        SC = 1
        for freq in self.frequences:
            if freq == '1min':
                ind = self.market_data_1min.groupby(level=1, sort=False).apply(QA.QA_indicator_MACD)
                # dfs.append(self.market_data_1min[['close', 'JC', 'SC']])
                last_bar = ind.xs(ind.index.levels[0][-1], level=0)
                if str(current_dt) == '2020-12-09 09:31:00':
                    print('current bar is ')
                    print(current_bar)
                    print('last_ind_bar is ')
                    print(last_bar)
                JC *= last_bar.JC
                SC *= last_bar.SC
            elif freq == '5min':
                ind = self.market_data_5min.groupby(level=1, sort=False).apply(QA.QA_indicator_MACD)
                last_bar = ind.xs(ind.index.levels[0][-1], level=0)
                JC *= last_bar.JC
                SC *= last_bar.SC
            elif freq == '15min':
                self._market_data_15min = self.market_data_15min.groupby(level=1, sort=False).apply(
                    QA.QA_indicator_MACD)
                # dfs.append(self.market_data_15min[['close', 'JC', 'SC']])
            elif freq == '30min':
                self._market_data_30min = self.market_data_30min.groupby(level=1, sort=False).apply(
                    QA.QA_indicator_MACD)
                # dfs.append(self.market_data_30min[['close', 'JC', 'SC']])
            elif freq == '60min':
                self._market_data_60min = self.market_data_60min.groupby(level=1, sort=False).apply(
                    QA.QA_indicator_MACD)
                # dfs.append(self.market_data_60min[['close', 'JC', 'SC']])
            elif freq == 'day':
                ind = self.market_data_day.groupby(level='code', sort=False).apply(QA.QA_indicator_MACD)
                last_bar = ind.xs(ind.index.levels[0][-1], level=0)
                JC *= last_bar.JC
                SC *= last_bar.SC
            elif freq == 'week':
                self._market_data_week = self.market_data_week.groupby(level=1, sort=False).apply(QA.QA_indicator_MACD)
                # dfs.append(self.market_data_week[['close', 'JC', 'SC']])
            elif freq == 'month':
                self._market_data_month = self.market_data_month.groupby(level=1, sort=False).apply(
                    QA.QA_indicator_MACD)
                # dfs.append(self.market_data_month[['close', 'JC', 'SC']])
            elif freq == 'year':
                self._market_data_year = self.market_data_year.groupby(level=1, sort=False).apply(QA.QA_indicator_MACD)

        print(current_dt, current_bar)

        # 按代码挨个判断买卖条件, 买入通过定额现金计算量， 卖出全部持仓
        for code in self.code:
            row = current_bar.loc[code]
            if JC.loc[code] > 0:
                amount = self.calc_vol_by_money(row.close, self.acc.available, self.once_trade_money)
                if self.enable_trade(amount * row.close):
                    self.send_order('BUY', 'OPEN', price=row.close, volume=amount, code=code)
            elif SC.loc[code] > 0:
                amount = self.acc.get_position(code).volume_long_his
                if amount > 0:
                    self.send_order('SELL', 'CLOSE', price=row.close, volume=-amount, code=self.code)

    def on_bar_backtest(self, bar):
        current_dt = bar.index.values[0][0]
        current_bar = bar.xs(current_dt, level=0)
        JC = 1
        SC = 1
        for freq in self.frequences:
            if freq == '1min':
                ind = self.market_data_1min.groupby(level=1, sort=False).apply(QA.QA_indicator_MACD)
                # dfs.append(self.market_data_1min[['close', 'JC', 'SC']])
                last_bar = ind.xs(ind.index.levels[0][-1], level=0)
                if str(current_dt) == '2020-12-09 09:31:00':
                    print('current bar is ')
                    print(current_bar)
                    print('last_ind_bar is ')
                    print(last_bar)
                JC *= last_bar.JC
                SC *= last_bar.SC
            elif freq == '5min':
                ind = self.market_data_5min.groupby(level=1, sort=False).apply(QA.QA_indicator_MACD)
                last_bar = ind.xs(ind.index.levels[0][-1], level=0)
                JC *= last_bar.JC
                SC *= last_bar.SC
            elif freq == '15min':
                self._market_data_15min = self.market_data_15min.groupby(level=1, sort=False).apply(QA.QA_indicator_MACD)
                # dfs.append(self.market_data_15min[['close', 'JC', 'SC']])
            elif freq == '30min':
                self._market_data_30min = self.market_data_30min.groupby(level=1, sort=False).apply(QA.QA_indicator_MACD)
                # dfs.append(self.market_data_30min[['close', 'JC', 'SC']])
            elif freq == '60min':
                self._market_data_60min = self.market_data_60min.groupby(level=1, sort=False).apply(QA.QA_indicator_MACD)
                # dfs.append(self.market_data_60min[['close', 'JC', 'SC']])
            elif freq == 'day':
                ind = self.market_data_day.groupby(level='code', sort=False).apply(QA.QA_indicator_MACD)
                last_bar = ind.xs(ind.index.levels[0][-1], level=0)
                JC *= last_bar.JC
                SC *= last_bar.SC
            elif freq == 'week':
                self._market_data_week = self.market_data_week.groupby(level=1, sort=False).apply(QA.QA_indicator_MACD)
                # dfs.append(self.market_data_week[['close', 'JC', 'SC']])
            elif freq == 'month':
                self._market_data_month = self.market_data_month.groupby(level=1, sort=False).apply(QA.QA_indicator_MACD)
                # dfs.append(self.market_data_month[['close', 'JC', 'SC']])
            elif freq == 'year':
                self._market_data_year = self.market_data_year.groupby(level=1, sort=False).apply(QA.QA_indicator_MACD)
                # dfs.append(self.market_data_year[['close', 'JC', 'SC']])

        # 将列表中的所有周期数据合并，金叉死叉为0或1的数据，所以使用乘法
        # df_x = dfs[0]
        # for d in dfs:
        #     df_x = df_x * d

        # sell all at end time
        if current_dt.day == self.end.day \
                and current_dt.month == self.end.month \
                and current_dt.year == self.end.year:
            if current_dt.hour == 15:  # min bar sell at 15:00
                for code in self.code:
                    self.sell(current_bar.loc[code], self.acc.get_position(code).volume_long, code)
        # if False:
        #     pass
        else:
            # 按代码挨个判断买卖条件, 买入通过定额现金计算量， 卖出全部持仓
            for code in self.code:
                row = current_bar.loc[code]
                if JC.loc[code] > 0:
                    amount = self.calc_vol_by_money(row.close, self.acc.available, self.once_trade_money)
                    if self.enable_trade(amount * row.close):
                        self.buy(row, amount, code)
                        if self.trade_start_time is None:
                            self.trade_start_time = bar.name[0].strftime("%Y-%m-%d %H:%M:%S")
                elif SC.loc[code] > 0:
                    amount = self.acc.get_position(code).volume_long_his
                    if amount > 0:
                        self.sell(row, amount, code)

    def calc_vol_by_money(self, price, cash, money):
        if money > cash:
            return 0
        else:
            vol = int(money / price / 100)
            if vol > 0:
                return vol
            else:
                return 0

    def calc_vol_by_pos(self, price, pos, cash, ratio):
        hold = price * (pos.volume_long)  # position money
        total = hold + cash
        hold_ratio = hold / total
        if ratio == 0:
            return -pos.volume_long_his
        elif hold_ratio > ratio:
            amount = int((hold - total * ratio) / price / 100) * 100
            if amount <= pos.volume_long_his:  # 卖出时需要考虑历史long仓位
                return -amount
            else:
                return 0
        elif hold_ratio < ratio:
            amount = int((total * ratio - hold) / price / 100.25) * 100
            return amount
        else:
            return 0

    def calc_vol(self, price, hold_vol, cash, ratio, towards=1):
        hold = price * hold_vol
        total = hold + cash
        hold_ratio = hold / total
        if hold_ratio != ratio:
            if towards != -1:  # BUY OR keep position grid
                return int((total * ratio - hold) / price / 100.25) * 100  # 100.25的意思是包含手续费
            else:
                return int((hold - total * ratio) / price / 100) * 100
        else:
            return 0

    def enable_trade(self, money):
        if abs(money) >= self.min_trade_money:
            return True
        else:
            return False

    def buy(self, bar, amount, code):
        order = self.acc.send_order(
            code=code,
            time=bar.name[0],
            amount=amount,
            towards=QA.ORDER_DIRECTION.BUY,
            price=bar.close,
            order_model=QA.ORDER_MODEL.CLOSE,
            amount_model=QA.AMOUNT_MODEL.BY_AMOUNT
        )
        if isinstance(order, bool):
            print('order is bool')
            return
        self.broker.receive_order(QA.QA_Event(order=order, market_data=bar))
        trade_mes = self.broker.query_orders(self.acc.account_cookie, 'filled')
        res = trade_mes.loc[order.account_cookie, order.realorder_id]
        order.trade(res.trade_id, res.trade_price, res.trade_amount, res.trade_time)
        if self.show_Trade_log:
            print(bar.name[0], 'BUY', bar.close, amount,
              int(self.acc.cash_available), int(self.acc.cash_available + self.acc.hold.get(code, 0) * bar.close))

    def sell(self, bar, amount, code):
        order = self.acc.send_order(
            code=code,
            time=bar.name[0],
            amount=amount,
            towards=QA.ORDER_DIRECTION.SELL,
            price=bar.close,
            order_model=QA.ORDER_MODEL.MARKET,
            amount_model=QA.AMOUNT_MODEL.BY_AMOUNT
        )
        # print
        if isinstance(order, bool):
            print('order is bool')
            return
        self.broker.receive_order(QA.QA_Event(order=order, market_data=bar))
        trade_mes = self.broker.query_orders(self.acc.account_cookie, 'filled')
        res = trade_mes.loc[order.account_cookie, order.realorder_id]
        order.trade(res.trade_id, res.trade_price, res.trade_amount, res.trade_time)
        if self.show_Trade_log:
            print(bar.name[0], 'SEL', 0.0, bar.close, amount, int(self.acc.cash_available),
                  int(self.acc.cash_available + self.acc.hold.get(code, 0) * bar.close))


